# coding:utf-8
import numpy as np
import cv2
import os
#-------------------------------
# 设置棋盘格大小  Checkerboard参数
#  scale这里scale就是边长。边长最终影响平移向量的大小
#-------------------------------
Checkerboard = (9,6)  #棋盘格内角点数
square_size = (10,10)   #棋盘格大小，单位mm
scale = square_size[0]
img_left_file = "./image/left"
img_right_file = "./image/right"

# 设置迭代终止条件
criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)                 # 单目的，亚像素函数使用
criteria_stereo = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 30, 0.001)           # 双目的。亚像素函数使用


# 设置 object points, 形式为 (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0) 
objp = np.zeros((Checkerboard[0] * Checkerboard[1], 3), np.float32)  # 我用的是6×9的棋盘格，可根据自己棋盘格自行修改相关参数
# objp[:, :2] = np.mgrid[0:Checkerboard[0], 0:Checkerboard[1]].T.reshape(-1, 2)

# 计算角点(棋盘格)的世界坐标系。 间隔为1的棋盘格的世界坐标系形式为 (0,0,0), (1,0,0), (2,0,0) ....,(6,5,0)  
objp[:, :2] = np.mgrid[0:(Checkerboard[0]-1)*scale:complex(0,Checkerboard[0]),
              0:(Checkerboard[1]-1)*scale:complex(0,Checkerboard[1])].T.reshape(-1, 2)
# obj_points = [] # 存储3D点
# print(objp)

# 用arrays存储所有图片的object points 和 image points
objpoints = []  # 棋盘格角点坐标(计算得到，绝对坐标，三维)3d points in real world space
imgpointsR = []  # 检测到的图片上的棋盘格角点坐标（检测得到，二维）2d points in image plane
imgpointsL = []


# imagL = glob.glob(r"./img/left/*.jpg")
# imagR = glob.glob(r"./img/right/*.jpg")   # 这里使用glob读取文件夹顺序会出问题，就采用土方法写死
# 本次实验采集里共计len(os.listdir(img_left_file))组待标定图片依次读入进行以下操作
print("一共%s对图像..." %len(os.listdir(img_left_file)))
for i in range( len(os.listdir(img_left_file)) ):
    # print(imagL[i],imagR[i])
    # print(img_left_file+"/"+str(i)+".jpg",img_right_file+"/"+str(i)+".jpg")
    ImaL = cv2.imread(img_left_file+"/left_"+str(i)+".jpg", 1)  # 右视图
    ImaR = cv2.imread(img_right_file+"/right_"+str(i)+".jpg", 1)  # 左视图
    grayL = cv2.cvtColor(ImaL, cv2.COLOR_BGR2GRAY)
    grayR = cv2.cvtColor(ImaR, cv2.COLOR_BGR2GRAY)
    retR, cornersR = cv2.findChessboardCorners(grayR, Checkerboard, None)  # 提取右图每一张图片的角点
    retL, cornersL = cv2.findChessboardCorners(grayL, Checkerboard, None)  # # 提取左图每一张图片的角点
    if (True == retR) & (True == retL):
        print("提取第%s的角点..."%i)
        objpoints.append(objp)
        cv2.cornerSubPix(grayL, cornersL, (11, 11), (-1, -1), criteria)  # 亚像素精确化，对粗提取的角点进行精确化
        cv2.cornerSubPix(grayR, cornersR, (11, 11), (-1, -1), criteria)  # 亚像素精确化，对粗提取的角点进行精确化
        imgpointsR.append(cornersR)
        imgpointsL.append(cornersL)

        cv2.drawChessboardCorners(ImaL, Checkerboard, cornersL, retL)
        cv2.drawChessboardCorners(ImaR, Checkerboard, cornersR, retR)

        # 把数据压到栈中,一起显示左右检测结果
        Image_merge = np.hstack((ImaL, ImaR))
        # 各参数依次是：图片，添加的文字，左上角坐标，字体，字体大小，颜色，字体粗细
        cv2.putText(Image_merge,"Left_img", (5, 40), 0, 1, (0, 0, 255), 2)
        cv2.putText(Image_merge, "Right_img", (5+ImaL.shape[1], 40), 0, 1, (0, 0, 255), 2)

        cv2.imshow('input img', Image_merge)
        # print("请敲击任意键，检测下一帧角点...")
        cv2.waitKey(0)
        cv2.destroyAllWindows()  # sx新增
    else:
        print("检测角点失败...,请检查图片路径或者棋盘格大小")
# cv2.destroyAllWindows()  # sx注释
# 相机的单双目标定、及校正
#   右侧相机单独标定
retR, mtxR, distR, rvecsR, tvecsR = cv2.calibrateCamera(objpoints, imgpointsR, grayR.shape[::-1], None, None)  
# 参数：objpoints角点绝对坐标，imgpointsR：检测到的角点坐标
# 返回：
# mtxR内参矩阵(3*3)，包含焦距f,主点坐标(u0,v0)，倾斜系数γ
# distR相机畸变，长度可为4(k1，k2，p1,和p2)、5(k1，k2，p1，p2和k3)、8(k1，k2，p1，p2，k3，k4，k5,和k6)，
# rvecsR：旋转矩阵。tvecsR：平移矩阵
# print("右相机内参：\n",mtxR, "右相机畸变：\n",distR)
# print("\n","-"*30,"\n")

#   获取新的相机矩阵后续传递给initUndistortRectifyMap，以用remap生成映射关系
hR, wR = grayR.shape[:2]
OmtxR, roiR = cv2.getOptimalNewCameraMatrix(mtxR, distR, (wR, hR), 1, (wR, hR))

#   左侧相机单独标定
retL, mtxL, distL, rvecsL, tvecsL = cv2.calibrateCamera(objpoints, imgpointsL, grayL.shape[::-1], None, None)
# print("左相机内参：\n",mtxL, "左相机畸变：\n",distL)
# print("\n","-"*30,"\n")
#   获取新的相机矩阵后续传递给initUndistortRectifyMap，以用remap生成映射关系
hL, wL = grayL.shape[:2]
OmtxL, roiL = cv2.getOptimalNewCameraMatrix(mtxL, distL, (wL, hL), 1, (wL, hL))

# 双目相机的标定
# 设置标志位为cv2.CALIB_FIX_INTRINSIC，这样就会固定输入的cameraMatrix和distCoeffs不变，只求解𝑅,𝑇,𝐸,𝐹
# flags = 1
# flags = cv2.CALIB_FIX_INTRINSIC
flags = cv2.CALIB_USE_INTRINSIC_GUESS

# 参数：存储标定角点在世界坐标系中的位,左边亚像素点，右边亚像素点，左内参+畸变，右内参+畸变,图像大小，
# 返回： 左相机内参矩阵MLS、左相机畸变dLS(与mtxL,distL值一样)、旋转矩阵R，平移向量T（R、T为两个相机坐标系之间的关系），本征矩阵E, 基础矩阵F，E、F用于极线几何，如求极线
retS, MLS, dLS, MRS, dRS, R, T, E, F = cv2.stereoCalibrate(objpoints, imgpointsL, imgpointsR,
                                                            mtxL, distL,
                                                           mtxR, distR,
                                                           grayR.shape[::-1],
                                                           criteria_stereo, flags)



# 利用stereoRectify()计算立体校正的映射矩阵
rectify_scale = 1  # 设置为0的话，对图片进行剪裁，设置为1则保留所有原图像像素
RL, RR, PL, PR, Q, roiL, roiR = cv2.stereoRectify(MLS, dLS, MRS, dRS, grayR.shape[::-1], R, T, rectify_scale, (0, 0))
# RL、RR：分别为左右相机消除畸变后的像平面投影到公共像平面的旋转矩阵
# PL、PR：分别为左右相机的投影矩阵，其作用是将世界坐标系的点转换到像素坐标系下(左相机光心为世界坐标系原点)
# roiL、roiR：用于标记remap图像中包含有效像素的矩形区域。
# Q矩阵为重投影矩阵，即矩阵Q可以把像素坐标系下的点转换到世界坐标系下。

# 利用initUndistortRectifyMap函数计算畸变矫正和立体校正的映射变换，实现极线对齐。
Left_Stereo_Map = cv2.initUndistortRectifyMap(MLS, dLS, RL, PL, grayR.shape[::-1], cv2.CV_16SC2)
Right_Stereo_Map = cv2.initUndistortRectifyMap(MRS, dRS, RR, PR, grayR.shape[::-1], cv2.CV_16SC2)

# 立体校正效果显示
for i in range(len(os.listdir(img_left_file))):  # 以第一对图片为例
    # print(img_left_file+"/left_"+str(i)+".jpg",img_right_file+"/right_"+str(i)+".jpg")
    frameR = cv2.imread(img_left_file+"/left_"+str(i)+".jpg", 1)
    frameL = cv2.imread(img_right_file+"/right_"+str(i)+".jpg", 1)

    Left_rectified = cv2.remap(frameL, Left_Stereo_Map[0], Left_Stereo_Map[1], cv2.INTER_LANCZOS4, cv2.BORDER_CONSTANT,
                               0)  # 使用remap函数完成映射
    # im_L = Image.fromarray(Left_rectified)  # numpy 转 image类

    Right_rectified = cv2.remap(frameR, Right_Stereo_Map[0], Right_Stereo_Map[1], cv2.INTER_LANCZOS4,
                                cv2.BORDER_CONSTANT, 0)
    # im_R = Image.fromarray(Right_rectified)  # numpy 转 image 类

    # 创建一个能同时并排放下两张图片的区域，后把两张图片依次粘贴进去
    # 建立输出图像
    height = max(Left_rectified.shape[0], Right_rectified.shape[0])
    width = Left_rectified.shape[1] + Right_rectified.shape[1]

    output = np.zeros((height, width, 3), dtype=np.uint8)

    output[0:Left_rectified.shape[0], 0:Left_rectified.shape[1]] = Left_rectified
    output[0:Right_rectified.shape[0], Right_rectified.shape[1]:] = Right_rectified

    line_interval = 40  # 直线间隔：40
    for k in range(height // line_interval):
        cv2.line(output, (0, line_interval * (k + 1)), (2 * width, line_interval * (k + 1)), (0, 255, 0), thickness=2,
                 lineType=cv2.LINE_AA)
    cv2.imshow("Stereo_rectified",output)
    print("按任意键显示下一帧对极几何图像...")
    cv2.waitKey()
cv2.destroyAllWindows()

# 打印标定结果
print("----------result of calibration-----------")
print("左相机内参：\n",mtxL, "\n左相机畸变：\n",distL)
print("右相机内参：\n",mtxR, "\n右相机畸变：\n",distR)
print("左右相机外参-旋转分量：\n",R,"\n左右相机外参-平移分量：\n",T)
print("-"*30,"\n")


# 保存结果
result = "cameraMatrixL：\n"+str(mtxL)+"\ndistCoeffL：\n"+str(distL)+ \
         "\ncameraMatrixR：\n"+str(mtxR)+"\ndistCoeffR：\n"+str(distR)+ \
         "\nR：\n"+str(R)+"\nT：\n"+str(T)
save_path = "result/stereo_calib.txt"
with open(save_path,'w') as f:
    f.write(result)
f.close()
print("save result successfully...")